# -*- coding: utf-8 -*-

from collections import OrderedDict
from functools import wraps
import collections

from flask_restful import marshal_with as orig_marshal_with
from flask_restful import marshal as orig_marshal
from flask_restful import unpack


def marshal(data, fields, envelope=None):
    # if isinstance(data, dict) and ('message' in data or 'data' in data):
    #     return data
    if isinstance(data, str):
        return {envelope: data} if envelope else data
    if isinstance(data, dict):
        new_data = {}
        for k, v in data.items():
            ret = marshal(v, fields, k)
            new_data.update(ret)
        return {envelope: new_data} if envelope else new_data
    # elif isinstance(data, collections.Iterator):
    #     return {'data': [orig_marshal(i, fields, envelope) for i in data]}
    # else:
    #     return {'data': orig_marshal(data, fields, envelope)}
    if envelope is None:
        envelope = 'data'
    if isinstance(data, collections.Iterator):
        return orig_marshal(list(data), fields, envelope)
    else:
        return orig_marshal(data, fields, envelope)


class marshal_with(orig_marshal_with):
    def __call__(self, f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            resp = f(*args, **kwargs)
            if isinstance(resp, tuple):
                data, code, headers = unpack(resp)
                return marshal(data, self.fields, self.envelope), code, headers
            else:
                return marshal(resp, self.fields, self.envelope)

        return wrapper

